# 导入 TensorFlow 库，用于加载和运行深度学习模型
import tensorflow as tf

# 导入 Pillow 库，用于图像处理
from PIL import Image

# 导入 NumPy 库，用于数值计算和数组操作
import numpy as np

# 导入 argparse 模块用于命令行参数解析
import argparse

# 加载预先训练好的字符识别模型
model = tf.keras.models.load_model('character_recognition_model.h5')

# 定义标签映射，将模型输出的索引转换为对应的字符
label_mapping = {0: 'V', 1: 'X', 2: 'O'}


def preprocess_image(image_path):
    # 打开指定路径的图片，并将其转换为灰度图
    img = Image.open(image_path).convert('L')
    # 调整图片大小为 28x28 像素，以匹配模型的输入要求
    img = img.resize((28, 28))
    # 将图片转换为 NumPy 数组，并将像素值归一化到 0-1 范围
    img_array = np.array(img) / 255.0
    # 增加一个维度，使图片数组形状变为 (1, 28, 28)，以匹配模型的输入形状
    img_array = np.expand_dims(img_array, axis=0)
    # 再次增加一个维度，使图片数组形状变为 (1, 28, 28, 1)，以匹配模型的输入形状
    img_array = np.expand_dims(img_array, axis=-1)
    return img_array


def predict_character(image_path):
    # 调用 preprocess_image 函数对图片进行预处理
    processed_image = preprocess_image(image_path)
    # 使用模型对预处理后的图片进行预测
    predictions = model.predict(processed_image)
    print(f"predictions: {predictions}")
    # 获取预测结果中概率最高的索引
    predicted_index = np.argmax(predictions)
    # 根据索引从 label_mapping 中获取对应的字符
    predicted_character = label_mapping[predicted_index]
    return predicted_character


def main():
    # 添加命令行参数解析
    parser = argparse.ArgumentParser(description="Character Recognition Inference")
    parser.add_argument("--image_path", type=str, required=True, help="Path to the input image")
    args = parser.parse_args()

    # 将 image_path 替换为从命令行参数中获取
    image_path = args.image_path

    # 调用 predict_character 函数进行预测，并获取结果
    result = predict_character(image_path)
    # 打印预测结果
    print(f"预测结果: {result}")

if __name__ == "__main__":
    main()